{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Investigate Wind Direction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import numpy as np\n",
    "import cartopy.crs as ccrs\n",
    "import geopandas as gpd\n",
    "import glob\n",
    "import moviepy.editor as mpy\n",
    "\n",
    "from palettable.colorbrewer.diverging import RdBu_11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "t:   0%|          | 2/8784 [00:00<09:46, 14.98it/s, now=None]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Building video movie.mp4.\n",
      "Moviepy - Writing video movie.mp4\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                  \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moviepy - Done !\n",
      "Moviepy - video ready movie.mp4\n"
     ]
    }
   ],
   "source": [
    "# file_list = sorted(glob.glob('movie_frame/*.png'))\n",
    "# clip = mpy.ImageSequenceClip(file_list, fps=30)\n",
    "# clip.write_videofile('movie.mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Alpine-3D netCDF file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Alpine-3D grids\n",
    "ds = xr.open_dataset(\"../output/grids/a3d_grids.nc\")\n",
    "\n",
    "# SNOWPACK topography \n",
    "dem = np.flipud(np.loadtxt(\"../input/modified_surface_grids/dem.asc\", skiprows=6))\n",
    "dem = xr.DataArray(dem, coords=[ds['northing'], ds['easting']], dims=['northing', 'easting'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trim grids\n",
    "ds = ds.isel(easting=slice(5, -5))\n",
    "ds = ds.isel(northing=slice(5, -5))\n",
    "\n",
    "dem = dem.isel(easting=slice(5, -5))\n",
    "dem = dem.isel(northing=slice(5, -5))\n",
    "dem = dem.values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot wind speed and direction from Alpine-3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "Calculates U and V components of wind from wind speed and direction. \n",
    "Wind direction is the meteorological wind direction (i.e. direction the wind is coming from in degrees)\n",
    "'''\n",
    "def calc_U_V(speed, direction):\n",
    "    U = -speed * np.sin(math.pi / 180 * direction)\n",
    "    V = -speed * np.cos(math.pi / 180 * direction)\n",
    "    return U, V\n",
    "\n",
    "# Get U and V components\n",
    "U, V = calc_U_V(ds['ws'], ds['dw'])\n",
    "VW_max = ds['ws'].max()\n",
    "VW_min = 0\n",
    "\n",
    "# Calculate\n",
    "dswe = 1000 * (ds['swe'] - ds['swe'].isel(time=0))\n",
    "\n",
    "# Get lat/lon\n",
    "x_snowpack = ds['easting']\n",
    "y_snowpack = ds['northing']\n",
    "\n",
    "\n",
    "# Make movie\n",
    "# Clear old images, gifs,and movies\n",
    "# !mkdir -p movie_frame\n",
    "# !rm -f movie.gif\n",
    "# !rm -f movie.mp4\n",
    "# !rm -f movie_frame/*\n",
    "\n",
    "# for time_step in range(0, len(ds['time']), 1):\n",
    "for time_step in range(5850, len(ds['time']), 1):\n",
    "    print(time_step)\n",
    "    plt.clf()\n",
    "    \n",
    "    # Plot map of mean wind\n",
    "    plt.figure(figsize=(30,15))\n",
    "\n",
    "    # DEM\n",
    "    contour_levels = np.linspace(dem.min(), dem.max(), 25)\n",
    "    contour = plt.contour(x_snowpack.values, y_snowpack.values, dem, contour_levels, linestyles='solid', colors='white')\n",
    "    plt.clabel(contour, fmt = '%.0f', inline = True)\n",
    "\n",
    "    # Delta SWE\n",
    "    colormap = RdBu_11.mpl_colormap # works\n",
    "    plt.pcolor(x_snowpack.values, y_snowpack.values, dswe.isel(time=time_step), \\\n",
    "               cmap = colormap, vmin=-np.abs(dswe).max(), vmax=np.abs(dswe).max())\n",
    "    plt.colorbar()\n",
    "\n",
    "    # Wind speed and direction (still need to normalize arrows so that length means the same thing for all time steps)\n",
    "#     spacing = 6\n",
    "    spacing = 2\n",
    "    northing_ind = np.arange(0, len(ds['northing']), spacing)\n",
    "    easting_ind = np.arange(0, len(ds['easting']), spacing)\n",
    "    plt.quiver(x_snowpack[easting_ind].values, y_snowpack[northing_ind].values, \\\n",
    "               U.isel(time=time_step)[northing_ind, easting_ind], V.isel(time=time_step)[northing_ind, easting_ind], scale=250)\n",
    "    plt.title(ds['time'].isel(time=time_step).values, fontsize=20)\n",
    "    \n",
    "    #Save Figure with image number zero padding \n",
    "    if time_step < 10: # 1 digit\n",
    "        plt.savefig(\"movie_frame/frame_000\" + str(time_step) + \".png\", dpi=100)\n",
    "    elif time_step < 100 and time_step > 9: # 2 digits\n",
    "        plt.savefig(\"movie_frame/frame_00\" + str(time_step) + \".png\", dpi=100)\n",
    "    elif time_step < 1000 and time_step > 99: # 3 digits\n",
    "        plt.savefig(\"movie_frame/frame_0\" + str(time_step) + \".png\", dpi=100)\n",
    "    else: # 4 digits\n",
    "        plt.savefig(\"movie_frame/frame_\" + str(time_step) + \".png\", dpi=100)\n",
    "    plt.close()\n",
    "    \n",
    "# Make a .mp4 movie and gif\n",
    "file_list = sorted(glob.glob('movie_frame/*.png'))\n",
    "clip = mpy.ImageSequenceClip(file_list, fps=30)\n",
    "clip.write_videofile('movie.mp4')\n",
    "# clip.write_gif('movie.gif')\n",
    "    \n",
    "\n",
    "# # Plot map of wind direction\n",
    "# plt.figure(figsize=(20,10))\n",
    "# plt.pcolor(x_snowpack.values, y_snowpack.values, ds['dw'].isel(time=time_step), cmap='twilight')\n",
    "# plt.colorbar()\n",
    "\n",
    "# # Plot horizontal transect of wind direction\n",
    "# plt.figure(figsize=(10,5))\n",
    "# ds['dw'][0,50,:].plot()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot a transect of wind speed and topography"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate transects\n",
    "ind = 25\n",
    "ws_transect = ds['ws'].isel(time=0)[:, ind]\n",
    "# dem_transect = dem[:, ind]\n",
    "dem_transect = np.flip(np.gradient(dem[:, ind]))\n",
    "\n",
    "# Plot both transects\n",
    "fig, ax1 = plt.subplots(figsize=(15,5))\n",
    "\n",
    "# Wind Speed\n",
    "color = 'tab:red'\n",
    "ax1.set_ylabel('Wind Speed', color=color)  \n",
    "ax1.plot(ws_transect, color=color)\n",
    "ax1.tick_params(axis='y', labelcolor=color)\n",
    "ax2 = ax1.twinx() \n",
    "\n",
    "# Elevation\n",
    "color = 'tab:blue'\n",
    "ax2.set_ylabel('Elevation Gradient', color=color)  \n",
    "ax2.plot(dem_transect, color=color)\n",
    "ax2.tick_params(axis='y', labelcolor=color)\n",
    "\n",
    "# Settings\n",
    "fig.tight_layout()  \n",
    "plt.title(\"Vertical Transect at Easting = \" + str(float(ws_transect['easting'].values)))\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plot MERRA-2 winds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define upper left and lower right A3D domain corners in epsg 3031\n",
    "ncol = 168\n",
    "nrow = 107\n",
    "left = -1542000 # done\n",
    "bottom = -109000 # done\n",
    "top = bottom + 1000 * (nrow -1)\n",
    "right = left + 1000 * (ncol -1)\n",
    "\n",
    "# Look at MERRA-2\n",
    "M2_U_path = \"/scratch/summit/erke2265/MERRA2/U10M_hourly_1980.nc\"\n",
    "M2_V_path = \"/scratch/summit/erke2265/MERRA2/V10M_hourly_1980.nc\"\n",
    "\n",
    "# Mean U and V components\n",
    "M2_U = xr.open_dataset(M2_U_path)\n",
    "M2_U = M2_U['U10M']\n",
    "M2_U_mean = M2_U[0,:,:].values\n",
    "# M2_U_mean = M2_U.mean(dim='time').values\n",
    "\n",
    "M2_V = xr.open_dataset(M2_V_path)\n",
    "M2_V = M2_V['V10M']\n",
    "M2_V_mean = M2_V[0,:,:].values\n",
    "# M2_V_mean = M2_V.mean(dim='time').values\n",
    "\n",
    "# Calculate mean wind speed.\n",
    "ws_mean = np.sqrt( np.power(M2_U_mean,2) + np.power(M2_V_mean,2))\n",
    "\n",
    "# Load coastlines\n",
    "df = gpd.read_file(\"/pl/active/nasa_smb/Data/ADD_Coastline_low_res_polygon.shp\")\n",
    "crs_epsg = ccrs.SouthPolarStereo()\n",
    "df_epsg = df.to_crs(epsg='3031')\n",
    "\n",
    "# Generate figure \n",
    "fig, axs = plt.subplots(1, 1, subplot_kw={'projection': crs_epsg},\n",
    "                        figsize=(15, 15))\n",
    "\n",
    "# Plot coastlines\n",
    "# axs.set_extent((-180, 180, -90, -60), ccrs.PlateCarree())\n",
    "axs.set_extent((-110, -70, -85, -70), ccrs.PlateCarree())\n",
    "axs.add_geometries(df_epsg['geometry'], crs=crs_epsg,\n",
    "                      facecolor='none', edgecolor='black')\n",
    "\n",
    "# Plot winds\n",
    "vector_crs = ccrs.PlateCarree()\n",
    "space = 1\n",
    "ws = plt.pcolor(M2_U['lon'].values, M2_U['lat'].values, ws_mean, transform=vector_crs, alpha=0.5)\n",
    "cbar = plt.colorbar(ws)\n",
    "streamlines = axs.streamplot(M2_U['lon'].values, M2_U['lat'].values, M2_U_mean, \\\n",
    "            M2_V_mean, color='k', transform=vector_crs, density=5)\n",
    "\n",
    "# Plot model domain\n",
    "plt.scatter(left, bottom, c='r', label=\"A3D Domain\")\n",
    "plt.scatter(left, top, c='r')\n",
    "plt.scatter(right, bottom, c='r')\n",
    "plt.scatter(right, top, c='r')\n",
    "plt.savefig(\"MERRA-2.pdf\", format='pdf', dpi=100)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "alpine3d",
   "language": "python",
   "name": "alpine3d"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
